package utils

import (
	"archive/zip"
	"io"
	"os"
	"path/filepath"
	"sync"
)

type ConcurrentZipper struct {
	Writer  *zip.Writer
	Counter int
	Mu      sync.Mutex
}

func NewConcurrentZipper(zipPath string, counter int) (*ConcurrentZipper, error) {
	ZipFile, err := os.Create(zipPath)
	if err != nil {
		return nil, err
	}

	return &ConcurrentZipper{
		Writer:  zip.NewWriter(ZipFile),
		Counter: counter,
	}, nil
}

func (z *ConcurrentZipper) Close() error {
	return z.Writer.Close()
}

func (z *ConcurrentZipper) AddFilesConcurrently(files []string) error {
	fileCh := make(chan string, len(files))
	errCh := make(chan error, len(files))
	var wg sync.WaitGroup

	// 启动 worker
	for i := 0; i < z.Counter; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for file := range fileCh {
				if err := z.AddFileToZip(file); err != nil {
					errCh <- err
				}
			}
		}()
	}

	// 发送文件
	for _, file := range files {
		fileCh <- file
	}
	close(fileCh)

	wg.Wait()
	close(errCh)

	for err := range errCh {
		return err
	}

	return nil
}

func (z *ConcurrentZipper) AddFileToZip(fileName string) error {
	z.Mu.Lock()
	defer z.Mu.Unlock()

	file, err := os.Open(fileName)
	if err != nil {
		return err
	}
	defer file.Close()
	info, err := file.Stat()
	if err != nil {
		return err
	}

	header, err := zip.FileInfoHeader(info)
	if err != nil {
		return err
	}
	header.Name = filepath.Base(fileName)
	header.Method = zip.Deflate

	writer, err := z.Writer.CreateHeader(header)
	if err != nil {
		return err
	}

	_, err = io.Copy(writer, file)
	return err
}

func (z *ConcurrentZipper) AddFilesSingly(filesMap map[string][]string) error {
	for k, v := range filesMap {
		for _, file := range v {
			if err := z.AddFileDirToZip(k, file); err != nil {
				return err
			}
		}
	}

	return nil
}

func (z *ConcurrentZipper) AddFileDirToZip(dir, fileName string) error {
	z.Mu.Lock()
	defer z.Mu.Unlock()

	file, err := os.Open(fileName)
	if err != nil {
		return err
	}
	defer file.Close()
	info, err := file.Stat()
	if err != nil {
		return err
	}

	header, err := zip.FileInfoHeader(info)
	if err != nil {
		return err
	}
	header.Name = filepath.Join(dir, filepath.Base(fileName))
	header.Method = zip.Deflate

	writer, err := z.Writer.CreateHeader(header)
	if err != nil {
		return err
	}

	_, err = io.Copy(writer, file)
	return err
}

func ZipFiles(zipPath string, files []string) error {
	zipFile, err := os.Create(zipPath)
	if err != nil {
		return err
	}
	defer zipFile.Close()

	zipWriter := zip.NewWriter(zipFile)
	defer zipWriter.Close()
	for _, file := range files {
		if err = addFileToZip(zipWriter, file); err != nil {
			return err
		}
	}
	return nil
}

func addFileToZip(zipWriter *zip.Writer, file string) error {
	fileToZip, err := os.Open(file)
	if err != nil {
		return err
	}
	defer fileToZip.Close()

	info, err := fileToZip.Stat()
	if err != nil {
		return err
	}

	header, err := zip.FileInfoHeader(info)
	if err != nil {
		return err
	}
	header.Name = filepath.Base(file)
	header.Method = zip.Deflate

	writer, err := zipWriter.CreateHeader(header)
	if err != nil {
		return err
	}

	_, err = io.Copy(writer, fileToZip)
	return err
}
